/*
    Copyright (C) 2010-2011  Anton Dedkov (Antonded). antonded@gmail.com

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <iostream>
#include <assert.h>
#include "compressor.h"

using std::ios_base;
using std::cout;
using std::endl;


//-----------------------------------------------------------------------------

Compressor::Compressor ():
tree_node_buffer_ (NULL),
buffer_index_ (0),
current_forest_ (),
file_read_buffer_ (NULL),
file_read_buffer_size_ (0)
{
    for (int i = 0; i < 257; ++i)
        profile_ [i] = NULL;
}

//-----------------------------------------------------------------------------

Compressor::~Compressor ()
{
    ClearAll ();
}


//-----------------------------------------------------------------------------

void Compressor::run ()
{
    
}

//-----------------------------------------------------------------------------

void Compressor::Compress (string in_file, string out_file)
{
    ReadToBuffer (in_file);
    
    Profile ();
    BuildTree ();

    fstream fout (out_file.c_str (), ios_base::out | ios_base::binary);
    assert (fout.is_open ());

    WriteProfile (fout);
    WriteCompressed (fout);

    fout.close ();

    ClearAll ();
}

//-----------------------------------------------------------------------------

void Compressor::Decompress (string in_file, string out_file)
{
    ReadToBuffer (in_file);
    ReadProfile();
    BuildTree ();

    fstream fout (out_file.c_str (), ios_base::out | ios_base::binary);
    assert (fout.is_open ());
    
    WriteDecompressed (fout);
    
    fout.close ();

    ClearAll ();
}

//-----------------------------------------------------------------------------

void Compressor::Dump () const
{
    cout << "\n-----------Compressor::Dump ()------------\n";
    cout << "profile_:\n";
    for (int i = 0; i < 257; ++i)
        if (profile_[i])
            cout << i << ":" << profile_[i]->freq << ' ';
    cout << endl;
    cout << "current_forest_:\n";
    TreeNode* temp = NULL;
    for (unsigned int i = 0; i < 257; ++i)
        if (profile_ [i])
        {
            cout << i << " ";
            temp = profile_ [i];
            while (temp->parent_)
            {
                if (temp == temp->parent_->child_ [0])
                    cout << 0;
                else
                    cout << 1;
            temp = temp->parent_;
            }
            cout << endl;
        }
}

//-----------------------------------------------------------------------------

void Compressor::Profile ()
{
    tree_node_buffer_ = new TreeNode [256 * 256 / 2];
    unsigned char it = 0;
    for (unsigned int i = 0; i < file_read_buffer_size_; ++i)
    {
        it = file_read_buffer_ [i];

        if (!(profile_ [it]))
        {
            profile_ [it] = tree_node_buffer_ + buffer_index_;
            buffer_index_++;

            profile_ [it]->freq = 0;
            profile_ [it]->parent_ = NULL;
            profile_ [it]->child_ [0] = profile_ [it]->child_ [1] = NULL;
            profile_ [it]->symbol = it;
        }

        profile_ [it]->freq++;
    }
    profile_ [eof_ - 1] = tree_node_buffer_ + buffer_index_;
    buffer_index_++;
    profile_ [eof_ - 1]->freq = 1;
    profile_ [eof_ - 1]->parent_ = NULL;
    profile_ [eof_ - 1]->child_ [0] = profile_ [eof_ - 1]->child_ [1] = NULL;
    profile_ [eof_ - 1]->symbol = eof_;
}

//-----------------------------------------------------------------------------

void Compressor::BuildTree ()
{
    for (unsigned int i = 0; i < 257; ++i)
        if (profile_ [i])
            current_forest_.push (profile_ [i]);

    while (current_forest_.size () > 1)
    {
        tree_node_buffer_ [buffer_index_].parent_ = NULL;
        tree_node_buffer_ [buffer_index_].freq = 0;
        tree_node_buffer_ [buffer_index_].symbol = 0;
        for (char i = 0; i < 2; ++i)
        {
            tree_node_buffer_ [buffer_index_].child_ [i] = current_forest_.top ();
            tree_node_buffer_ [buffer_index_].freq += current_forest_.top ()->freq;
            current_forest_.top ()->parent_ = tree_node_buffer_ + buffer_index_;
            current_forest_.pop ();
        }
    current_forest_.push (tree_node_buffer_ + buffer_index_);
    buffer_index_++;
    }
}

//-----------------------------------------------------------------------------

bool Compressor::TreeNodeCompare::operator() (const Compressor::TreeNode* rnode1,
                                           const Compressor::TreeNode* rnode2 )
{
    return rnode1->freq >= rnode2->freq;
}

//-----------------------------------------------------------------------------

void Compressor::WriteProfile (std::fstream& fout)
{
    unsigned long long int null_freq = 0;
    for (int i = 0; i < 256; ++i)
        if (profile_ [i])
            fout.write ((char*) &(profile_ [i]->freq), sizeof (long long int));
        else
            fout.write ((char*) &null_freq, sizeof (long long int));
}

//-----------------------------------------------------------------------------

void Compressor::WriteCompressed (std::fstream& fout)
{
    unsigned char bit_count = 0;
    unsigned char bit_num = 0;
    unsigned long long int bit_size = 0;
    const int buffer_size = 32;
    char buffer [buffer_size] = {0};
    char file_buffer [file_buffer_size_] = {0};
    TreeNode* temp = NULL;
    int file_read_buffer_i = 0;
    for (int i = 0; i <= file_read_buffer_size_; ++i)
    {
        if (i == file_read_buffer_size_)
            file_read_buffer_i = 256;
        else
            file_read_buffer_i = file_read_buffer_ [i];
        if (profile_ [file_read_buffer_i])
        {
            for (int j = 0; j < buffer_size; ++j)
                buffer [j] = 0;
            bit_count = 0;
            temp = profile_ [file_read_buffer_i];
            while (temp->parent_)
            {
                if (temp == temp->parent_->child_ [1])
                {
                    //write bit to next position
                    (*(buffer + bit_count / 8)) |= (1 << (7 - bit_count % 8));
                }
                bit_count++;
                temp = temp->parent_;
            }
            for (bit_num = 0; bit_num < bit_count; bit_num++, bit_size++)
            {
                if (bit_size >= file_buffer_size_ * 8 )
                {
                    fout.write (file_buffer, (bit_size + 7) / 8);
                    bit_size = 0;
                    for (int j = 0; j < file_buffer_size_; ++j)
                        file_buffer [j] = 0;
                }

                //write bits from buffer to file_buffer in reverse sequence
                //bit_size is current bit index in file_buffer
                //bit_count is count of bits in current symbol
                //bit_num is the bit number in symbol
                (*(file_buffer + bit_size / 8)) |=
                        (((buffer [(bit_count - bit_num - 1) / 8] >>
                        (7 - (bit_count % 8) + bit_num % 8 + 1 ) % 8) & 1) <<
                                                            (7 - bit_size % 8));
            }
        }
    }
    fout.write (file_buffer, (bit_size + 7) / 8);

}

//-----------------------------------------------------------------------------

void Compressor::WriteDecompressed (fstream& fout)
{
    char file_buffer [file_buffer_size_] = {0};
    unsigned int buffer_symbol_num = 0;
    TreeNode* temp = current_forest_.top ();
    for (int i = 256 * sizeof (long long int); i < file_read_buffer_size_; ++i)
        for (int bit_num = 0; bit_num < 8; ++bit_num)
        {
            if (temp->child_ [0] || temp->child_ [1])
                temp = temp->child_ [(file_read_buffer_ [i] >> (7 - bit_num)) & 1];
            else
            {
                if (temp->symbol == eof_)
                    break;
                file_buffer [buffer_symbol_num] = temp->symbol;
                buffer_symbol_num++;
                if (buffer_symbol_num >= file_buffer_size_)
                {
                    fout.write (file_buffer, buffer_symbol_num);
                    for (int j = 0; j < file_buffer_size_; ++j)
                        file_buffer [j] = 0;
                    buffer_symbol_num = 0;
                }
                temp = current_forest_.top ()->
                          child_ [(file_read_buffer_ [i] >> (7 - bit_num)) & 1];
            }
        }
    fout.write (file_buffer, buffer_symbol_num);
}


//-----------------------------------------------------------------------------

void Compressor::ReadToBuffer (string file)
{
    fstream fin (file.c_str (), ios_base::in | ios_base::binary);
    assert (fin.is_open ());
    
    fin.seekg (0, ios_base::end);
    file_read_buffer_size_ = fin.tellg ();
    file_read_buffer_ = new unsigned char [file_read_buffer_size_];
    fin.seekg (0, ios_base::beg);
    fin.read (reinterpret_cast <char*> (file_read_buffer_), file_read_buffer_size_);
    fin.close ();
}

//-----------------------------------------------------------------------------

void Compressor::ClearAll ()
{
    if (file_read_buffer_)
        delete [] file_read_buffer_;
    if (tree_node_buffer_)
        delete [] tree_node_buffer_;

    buffer_index_ = 0;
    while (current_forest_.size ())
        current_forest_.pop ();
    file_read_buffer_ = NULL;
    file_read_buffer_size_ = 0;
    for (int i = 0; i < 257; ++i)
        profile_ [i] = 0;
    tree_node_buffer_ = NULL;
}

//-----------------------------------------------------------------------------

void Compressor::ReadProfile()
{
    assert (file_read_buffer_size_ >= 256 * sizeof (long long int));
    tree_node_buffer_ = new TreeNode [256 * 256 / 2];

    for (unsigned int i = 0; i < 256; ++i)
    {
        if (!(*((long long int*) (file_read_buffer_ + sizeof (long long int) * i))))
            continue;
        
        profile_ [i] = tree_node_buffer_ + buffer_index_;
        buffer_index_++;
            
        profile_ [i]->freq =
           (*((long long int*) (file_read_buffer_ + sizeof (long long int) * i)));
        profile_ [i]->parent_ = NULL;

        profile_ [i]->child_ [0] = profile_ [i]->child_ [1] = NULL;
        profile_ [i]->symbol = i;
    }
    profile_ [eof_ - 1] = tree_node_buffer_ + buffer_index_;
    profile_ [eof_ - 1]->freq = 1;
    profile_ [eof_ - 1]->parent_ = NULL;
    profile_ [eof_ - 1]->child_ [0] = profile_ [eof_ - 1]->child_ [1] = NULL;
    profile_ [eof_ - 1]->symbol = eof_;

    buffer_index_++;
    
}
